package com.study.repository.support;

import com.study.anno.Collection;
import com.study.dto.SimilarityResultDto;
import com.study.parse.ParseParam;
import com.study.repository.MilvusRepository;
import io.milvus.client.MilvusServiceClient;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.grpc.DataType;
import io.milvus.grpc.FieldData;
import io.milvus.grpc.ListDatabasesResponse;
import io.milvus.grpc.SearchResults;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.*;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import jakarta.annotation.Resource;
import org.springframework.stereotype.Repository;

import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;

@Repository
public class SimpleMilvusRepository<T, ID> implements MilvusRepository<T, ID> {

    private static final String ID_MUST_NOT_BE_NULL = "The given id must not be null!";

    private Class<T> entityType;

    @Resource
    private ParseParam parseParam;

    @Resource
    private MilvusServiceClient milvusClient;

    final String databaseName = "test";

    @Override
    public List<SimilarityResultDto<T>> searchVector(Class<T> clazz, int SEARCH_K, float[] vector, String searchStr) {

//        Type superClass = getClass().getGenericSuperclass();
//        if (superClass instanceof ParameterizedType) {
//            ParameterizedType parameterizedType = (ParameterizedType) superClass;
//            this.entityType = (Class<T>) parameterizedType.getActualTypeArguments()[0];
//        }

        R<ListDatabasesResponse> databaseResponse = milvusClient.listDatabases();
        if (!databaseResponse.getData().getDbNamesList().contains(databaseName)) {
            milvusClient.createDatabase(CreateDatabaseParam.newBuilder().withDatabaseName(databaseName).build());
        }

        Collection collection = clazz.getDeclaredAnnotation(Collection.class);
        String collectionName = collection.collectionName();

        R<Boolean> result = milvusClient.hasCollection(HasCollectionParam.newBuilder()
                .withDatabaseName(databaseName).withCollectionName(collectionName).build());
        if (result.getData()) {
            R<RpcStatus> status = milvusClient.loadCollection(LoadCollectionParam.newBuilder()
                    .withDatabaseName(databaseName)
                    .withCollectionName(collectionName)
                    .build());
            System.out.println("获取表结果:" + status.getData());

            status = milvusClient.dropCollection(DropCollectionParam.newBuilder()
                    .withDatabaseName(databaseName)
                    .withCollectionName(collectionName).build());
            System.out.println("删除表结果:" + status.getData());
        }

        FieldType id = FieldType.newBuilder()
                .withName("id")
                .withDataType(DataType.Int64)
                .withPrimaryKey(true)
                .withAutoID(true)
                .build();

        FieldType name = FieldType.newBuilder()
                .withName("name")
                .withDataType(DataType.VarChar)
                .withMaxLength(32)
                .withPrimaryKey(false)
                .withAutoID(false)
                .build();

        FieldType features = FieldType.newBuilder()
                .withName("features")
                .withDataType(DataType.FloatVector)
                .withPrimaryKey(false)
                .withAutoID(false)
                .withDimension(32)
                .build();


        R<RpcStatus> status = milvusClient.createCollection(
                CreateCollectionParam.newBuilder()
                        .withCollectionName(collectionName)
                        .withDatabaseName(databaseName)
                        .withSchema(CollectionSchemaParam.newBuilder()
                                .withFieldTypes(List.of(id, name, features))
                                .build())
                        .build());
        System.out.println("创建表结果:" + status);

        // 创建索引
        R<RpcStatus> r = milvusClient.createIndex(CreateIndexParam.newBuilder()
                .withDatabaseName(databaseName)
                .withCollectionName(collectionName)
                .withFieldName("features")
                .withIndexType(IndexType.IVF_FLAT)
                .withMetricType(MetricType.L2)
                .withExtraParam("{\"nlist\":100}")
                .build());
        System.out.println("创建索引结果:" + r);

        //查询之前,需要先loadCollection开启数据自动加载到内存
        milvusClient.loadCollection(LoadCollectionParam.newBuilder()
                        .withDatabaseName(databaseName)
                .withCollectionName(collectionName)
                .build());

        List<String> outFields = getQueryParam(clazz);
        List<List<Float>> search_vectors = List.of(arrayToList(vector));
        SearchParam searchParam = SearchParam.newBuilder()
                .withDatabaseName(databaseName)
                .withCollectionName(collectionName)
                .withConsistencyLevel(ConsistencyLevelEnum.STRONG)
                .withMetricType(MetricType.COSINE)
                .withOutFields(outFields)
                .withTopK(SEARCH_K)
                .withVectors(search_vectors)
                .withVectorFieldName("features")
                .withParams(searchStr)
                .build();
        R<SearchResults> respSearch = milvusClient.search(searchParam);
        return getResult(respSearch, clazz, SEARCH_K);
    }

    private <T> List<String> getQueryParam(Class<T> clazz) {
        if (Objects.nonNull(clazz)) {
            return Arrays.stream(clazz.getDeclaredFields()).map(Field::getName).toList();
        }
        throw new RuntimeException("class not define");
    }

    private List<Float> arrayToList(float[] array) {
        List<Float> floatList = new ArrayList<>();
        for (float value : array) {
            floatList.add(value);
        }
        return floatList;
    }

    private <U> List<SimilarityResultDto<U>> getResult(R<SearchResults> resultsR, Class<U> clazz, int SEARCH_K) {
        if (resultsR == null) {
            throw new IllegalArgumentException("Results cannot be null");
        }
        SearchResults searchResults = resultsR.getData();
        if (Objects.isNull(searchResults)) {
            return new ArrayList<>();
        }

        try {
            List<FieldData> fieldDataList = searchResults.getResults().getFieldsDataList();
            if (fieldDataList.isEmpty()) {
               return new ArrayList<>();
            }

            List<SimilarityResultDto<U>> ret = new ArrayList<>();

            List<String> queryField = fieldDataList.stream()
                    .map(FieldData::getFieldName)
                    .toList();
            // 使用Optional避免直接空指针检查
            U ins = clazz.getDeclaredConstructor().newInstance();
            Field[] declaredFields = clazz.getDeclaredFields();
            for (int i = 0; i < fieldDataList.size(); i++) {
                SimilarityResultDto<U> dto = new SimilarityResultDto<>();
                for (Field field : declaredFields) {
                    int index = queryField.indexOf(field.getName());
                    if (index != -1) {
                        processField(ins, field, fieldDataList.get(index), i);
                        dto.setData(ins);
                    }
                }
                dto.setSimilarity(resultsR.getData().getResults().getScores(i));
                ret.add(dto);
            }
            return ret;
        } catch (InstantiationException | IllegalAccessException | InvocationTargetException |
                 NoSuchMethodException e) {
            throw new RuntimeException("反射操作失败: " + e.getMessage(), e);
        }
    }

    private <U> void processField(U ins, Field field, FieldData fieldData, int index)
            throws IllegalAccessException, NoSuchMethodException, InvocationTargetException {
        field.setAccessible(true);
        parseParam.resolveParamBind(ins, field, fieldData, index);
    }

}
